import numpy as np

from sklearn import datasets

iris = datasets.load_iris()
dataset = [(iris.data[i][None, ...], iris.target[i]) for i in range(len(iris.target))]

INPUT_DIM = 4  # входные параметры
OUTPUT_DIM = 3  # ответы нейронной сети

H_DIM = 5  # Количество нейронов в слое


# функции

def relu(x):  # функция активации 1 скрытого слоя
    return np.maximum(0, x)


def softmax(x):  # функция активации выходного слоя
    out = np.exp(x)
    return out / np.sum(out)


def sparse_cross_entropy(a, b):  # ошибка
    # измерение расстояния между двумя вероятносными измерениями лучше использовать функцию кросс-энтропии
    # функция кросс-энтропия -Σy_i*log(z_i)
    return -np.log(a[0, b])


def to_full(a, num_class): # правильный ответ записываем в вектор чтобы размерности совпали
    y_full = np.zeros((1, num_class))
    y_full[0, a] = 1
    return y_full


def relu_deriv(t):  # производная от функция ReLU
    return (t >= 0).astype(float)


# два полносвязных слоя
# объявление весов и bias
W1 = np.random.randn(INPUT_DIM, H_DIM)  # размерность (INPUT_DIM, H_DIM)
b1 = np.random.randn(1, H_DIM)
W2 = np.random.randn(H_DIM, OUTPUT_DIM)
b2 = np.random.randn(1, OUTPUT_DIM)

ALPHA = 0.01  # скорость обучения
EPOCH = 200  # количество эпох
loss_arr = []

for epoh in range(EPOCH):
    for i in range(len(dataset)):
        x, y = dataset[i]

        # Forward

        # @ - матричное умножение

        t1 = x @ W1 + b1 # вычисление 1 скрытого слоя
        h1 = relu(t1) # функция активации промежуточного выходного слоя
        t2 = h1 @ W2 + b2 # вычисление 2 скрытого слоя
        z = softmax(t2) # функция активации выходного слоя

        E = sparse_cross_entropy(z, y)  # ошибка

        # Backward

        y_full = to_full(y, OUTPUT_DIM) # правильный ответ записываем в ветор той же размерности что и output_dim

        ''' вычисление градиентов 
        dE/dt1 <- dE/dh1 <- dE/dt2 <- E
        из dE/dt2 можем получить dE/dW2 и dE/db2
        из dE/dt1 можем получить dE/dW1 и dE/db1 '''

        dE_dt2 = z - y_full
        dE_dW2 = h1.T @ dE_dt2
        dE_b2 = dE_dt2

        dE_h1 = dE_dt2 @ W2.T

        dE_dt1 = dE_h1 * relu_deriv(t1)
        dE_dW1 = x.T @ dE_dt1
        dE_db1 = dE_dt1

        # Update

        W1 -= ALPHA * dE_dW1
        b1 -= ALPHA * dE_db1
        W2 -= ALPHA * dE_dW2
        b2 -= ALPHA * dE_b2

        loss_arr.append(E)


def predict(x):
    t1 = x @ W1 + b1  # @ - матричное умножение
    h1 = relu(t1)
    t2 = h1 @ W2 + b2
    z = softmax(t2)
    return z


def calc_accuracy():
    correct = 0
    for x, y in dataset:
        z = predict(x)
        y_pred = np.argmax(z)
        if y_pred == y:
            correct += 1
    acc = correct / len(dataset)
    return acc


accuracy = calc_accuracy()
print("Accuracy: ", accuracy)

import matplotlib.pyplot as plt

plt.plot(loss_arr)
plt.show()

x_new = [4.9, 3. , 1.4, 0.2]
class_iris = ['setosa', 'versicolor', 'virginica']
print(class_iris[np.argmax(predict(x_new))])